{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8ebf73ae-a632-4d09-b941-6739e35b760d",
   "metadata": {},
   "source": [
    "# An Introduction to LlamaIndex Query Pipelines\n",
    "\n",
    "## Overview\n",
    "LlamaIndex provides a declarative query API that allows you to chain together different modules in order to orchestrate simple-to-advanced workflows over your data.\n",
    "\n",
    "This is centered around our `QueryPipeline` abstraction. Load in a variety of modules (from LLMs to prompts to retrievers to other pipelines), connect them all together into a sequential chain or DAG, and run it end2end.\n",
    "\n",
    "**NOTE**: You can orchestrate all these workflows without the declarative pipeline abstraction (by using the modules imperatively and writing your own functions). So what are the advantages of `QueryPipeline`? \n",
    "\n",
    "- Express common workflows with fewer lines of code/boilerplate\n",
    "- Greater readability\n",
    "- Greater parity / better integration points with common low-code / no-code solutions (e.g. LangFlow)\n",
    "- [In the future] A declarative interface allows easy serializability of pipeline components, providing portability of pipelines/easier deployment to different systems.\n",
    "\n",
    "## Cookbook\n",
    "\n",
    "In this cookbook we give you an introduction to our `QueryPipeline` interface and show you some basic workflows you can tackle.\n",
    "\n",
    "- Chain together prompt and LLM\n",
    "- Chain together query rewriting (prompt + LLM) with retrieval\n",
    "- Chain together a full RAG query pipeline (query rewriting, retrieval, reranking, response synthesis)\n",
    "- Setting up a custom query component"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1144b2d4-adbc-44da-8c12-bdb5fe4b18bb",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Here we setup some data + indexes (from PG's essay) that we'll be using in the rest of the cookbook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adc82744-965a-4d79-b357-faf3de7ba2f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🌍 To view the Phoenix app in your browser, visit http://127.0.0.1:6006/\n",
      "📺 To view the Phoenix app in a notebook, run `px.active_session().view()`\n",
      "📖 For more information on how to use Phoenix, check out https://docs.arize.com/phoenix\n"
     ]
    }
   ],
   "source": [
    "# setup Arize Phoenix for logging/observability\n",
    "import phoenix as px\n",
    "\n",
    "px.launch_app()\n",
    "import llama_index\n",
    "\n",
    "llama_index.set_global_handler(\"arize_phoenix\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40fcb621-0894-457e-b602-dbf5fb9134ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.query_pipeline.query import QueryPipeline\n",
    "from llama_index.llms import OpenAI\n",
    "from llama_index.prompts import PromptTemplate\n",
    "from llama_index import (\n",
    "    VectorStoreIndex,\n",
    "    ServiceContext,\n",
    "    SimpleDirectoryReader,\n",
    "    load_index_from_storage,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2009af96-59e3-4d14-8272-382203c8b8a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "reader = SimpleDirectoryReader(\"../data/paul_graham\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c55d390b-38ca-4176-8cdb-8c2a0af1add8",
   "metadata": {},
   "outputs": [],
   "source": [
    "docs = reader.load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "780ddc81-0783-4fa5-ade0-60700c918011",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from llama_index.storage import StorageContext\n",
    "\n",
    "if not os.path.exists(\"storage\"):\n",
    "    index = VectorStoreIndex.from_documents(docs)\n",
    "    # save index to disk\n",
    "    index.set_index_id(\"vector_index\")\n",
    "    index.storage_context.persist(\"./storage\")\n",
    "else:\n",
    "    # rebuild storage context\n",
    "    storage_context = StorageContext.from_defaults(persist_dir=\"storage\")\n",
    "    # load index\n",
    "    index = load_index_from_storage(storage_context, index_id=\"vector_index\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5c59b5c-9b18-4dfa-97ef-e39b8069b73c",
   "metadata": {},
   "source": [
    "## 1. Chain Together Prompt and LLM \n",
    "\n",
    "In this section we show a super simple workflow of chaining together a prompt with LLM.\n",
    "\n",
    "We simply define `chain` on initialization. This is a special case of a query pipeline where the components are purely sequential, and we automatically convert outputs into the right format for the next inputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb233a0f-2993-4780-a241-6a2299047598",
   "metadata": {},
   "outputs": [],
   "source": [
    "# try chaining basic prompts\n",
    "prompt_str = \"Please generate related movies to {movie_name}\"\n",
    "prompt_tmpl = PromptTemplate(prompt_str)\n",
    "llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
    "\n",
    "p = QueryPipeline(chain=[prompt_tmpl, llm], verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b26c0c6-b886-42a6-b524-b19b18d1c01c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;155;135;227m> Running module 8d886c0b-5091-4650-af32-48dc62e6cdf1 with input: \n",
      "movie_name: The Departed\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module 9f5fe0b1-bb02-4b04-a11a-ad4f5eb33916 with input: \n",
      "messages: Please generate related movies to The Departed\n",
      "\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "output = p.run(movie_name=\"The Departed\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b987554f-4214-409f-97a8-40d2b13c25ee",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "assistant: 1. Infernal Affairs (2002) - This Hong Kong crime thriller is the original film on which The Departed is based. It follows a similar storyline of undercover cops infiltrating a criminal organization.\n",
      "\n",
      "2. Internal Affairs (1990) - This American crime thriller, starring Richard Gere and Andy Garcia, revolves around a corrupt cop and an internal affairs officer determined to expose him.\n",
      "\n",
      "3. The Town (2010) - Directed by and starring Ben Affleck, this crime drama follows a group of bank robbers in Boston who find themselves in a dangerous situation when they take a hostage during a heist.\n",
      "\n",
      "4. Mystic River (2003) - Directed by Clint Eastwood, this psychological crime drama features an ensemble cast including Sean Penn, Tim Robbins, and Kevin Bacon. It explores the aftermath of a childhood trauma and its impact on three friends.\n",
      "\n",
      "5. The Wire (TV Series, 2002-2008) - Although not a movie, this critically acclaimed TV series created by David Simon delves into the interconnected lives of law enforcement, drug dealers, and residents of Baltimore. It offers a complex and realistic portrayal of crime and corruption.\n",
      "\n",
      "6. Training Day (2001) - Starring Denzel Washington and Ethan Hawke, this crime thriller follows a rookie cop who is assigned to an experienced but corrupt detective for his first day on the job.\n",
      "\n",
      "7. American Gangster (2007) - Based on a true story, this crime film directed by Ridley Scott stars Denzel Washington as a Harlem drug lord and Russell Crowe as the detective determined to bring him down.\n",
      "\n",
      "8. The Departed (2006) (Remake) - Although not a separate movie, it is worth mentioning that The Departed is a remake of the original Hong Kong film Infernal Affairs. Both films have their own unique qualities and are worth watching.\n",
      "\n",
      "9. Donnie Brasco (1997) - Starring Johnny Depp and Al Pacino, this crime drama is based on the true story of an FBI agent who infiltrates the mob and forms a close bond with a mafia hitman.\n",
      "\n",
      "10. City of God (2002) - This Brazilian crime drama depicts the brutal and violent world of organized crime in the slums of Rio de Janeiro. It follows the lives of two young boys who take different paths, one becoming a photographer and the other a drug dealer.\n"
     ]
    }
   ],
   "source": [
    "print(str(output))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f44bac2-fff8-4578-8179-2cf68d075429",
   "metadata": {},
   "source": [
    "### Try Output Parsing\n",
    "\n",
    "Let's parse the outputs into a structured Pydantic object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69fc9aa6-d74e-4d02-8101-83ee03d68d52",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "from pydantic import BaseModel, Field\n",
    "from llama_index.output_parsers import PydanticOutputParser\n",
    "\n",
    "\n",
    "class Movie(BaseModel):\n",
    "    \"\"\"Object representing a single movie.\"\"\"\n",
    "\n",
    "    name: str = Field(..., description=\"Name of the movie.\")\n",
    "    year: int = Field(..., description=\"Year of the movie.\")\n",
    "\n",
    "\n",
    "class Movies(BaseModel):\n",
    "    \"\"\"Object representing a list of movies.\"\"\"\n",
    "\n",
    "    movies: List[Movie] = Field(..., description=\"List of movies.\")\n",
    "\n",
    "\n",
    "llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
    "output_parser = PydanticOutputParser(Movies)\n",
    "json_prompt_str = \"\"\"\\\n",
    "Please generate related movies to {movie_name}. Output with the following JSON format: \n",
    "\"\"\"\n",
    "json_prompt_str = output_parser.format(json_prompt_str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "243b55ca-a7bd-43e4-a837-20e29b3bebed",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;155;135;227m> Running module 246c804a-47a5-4cdd-a83b-90064d82740b with input: \n",
      "movie_name: Toy Story\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module 51d2c847-6a34-4a0d-841a-68cf009f8ded with input: \n",
      "messages: Please generate related movies to Toy Story. Output with the following JSON format: \n",
      "\n",
      "\n",
      "\n",
      "Here's a JSON schema to follow:\n",
      "{\"title\": \"Movies\", \"description\": \"Object representing a list of movies.\", \"typ...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module 1589c44b-a136-4093-abd7-286cb5d19d47 with input: \n",
      "input: assistant: {\n",
      "  \"movies\": [\n",
      "    {\n",
      "      \"name\": \"Finding Nemo\",\n",
      "      \"year\": 2003\n",
      "    },\n",
      "    {\n",
      "      \"name\": \"Cars\",\n",
      "      \"year\": 2006\n",
      "    },\n",
      "    {\n",
      "      \"name\": \"Up\",\n",
      "      \"year\": 2009\n",
      "    },\n",
      "    {...\n",
      "\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "# add JSON spec to prompt template\n",
    "json_prompt_tmpl = PromptTemplate(json_prompt_str)\n",
    "\n",
    "p = QueryPipeline(chain=[json_prompt_tmpl, llm, output_parser], verbose=True)\n",
    "output = p.run(movie_name=\"Toy Story\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b15bef7-399b-4dcc-94d6-6a4ea0066a41",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Movies(movies=[Movie(name='Finding Nemo', year=2003), Movie(name='Cars', year=2006), Movie(name='Up', year=2009), Movie(name='Inside Out', year=2015), Movie(name='Coco', year=2017)])"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fdfacf4-fb27-423a-b99c-37a84a1e2fbc",
   "metadata": {},
   "source": [
    "### Streaming Support\n",
    "\n",
    "The query pipelines have LLM streaming support (simply do `as_query_component(streaming=True)`). Intermediate outputs will get autoconverted, and the final output can be a streaming output. Here's some examples. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab92e2e2-9f66-47d3-9dbf-db6c7587e893",
   "metadata": {},
   "source": [
    "**1. Chain multiple Prompts with Streaming**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ba999ee-2281-46e5-9538-61297d73fe2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_str = \"Please generate related movies to {movie_name}\"\n",
    "prompt_tmpl = PromptTemplate(prompt_str)\n",
    "# let's add some subsequent prompts for fun\n",
    "prompt_str2 = \"\"\"\\\n",
    "Here's some text:\n",
    "\n",
    "{text}\n",
    "\n",
    "Can you rewrite this with a summary of each movie?\n",
    "\"\"\"\n",
    "prompt_tmpl2 = PromptTemplate(prompt_str2)\n",
    "llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
    "llm_c = llm.as_query_component(streaming=True)\n",
    "\n",
    "p = QueryPipeline(\n",
    "    chain=[prompt_tmpl, llm_c, prompt_tmpl2, llm_c], verbose=True\n",
    ")\n",
    "# p = QueryPipeline(chain=[prompt_tmpl, llm_c], verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c00485c-b5d5-4504-80c0-c8614e51523d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;155;135;227m> Running module c6f3bd13-5f78-4d37-9a00-4b23d3172619 with input: \n",
      "movie_name: The Dark Knight\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module d3c6bdb9-e891-428f-ae75-ef4181ba9dbd with input: \n",
      "messages: Please generate related movies to The Dark Knight\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module 8c84b2be-a338-4018-9cee-cc9d94a88b79 with input: \n",
      "text: <generator object llm_chat_callback.<locals>.wrap.<locals>.wrapped_llm_chat.<locals>.wrapped_gen at 0x2978e4e40>\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module d6203b7b-57bd-499a-aad8-3dfa7de093e5 with input: \n",
      "messages: Here's some text:\n",
      "\n",
      "1. Batman Begins (2005)\n",
      "2. The Dark Knight Rises (2012)\n",
      "3. Batman v Superman: Dawn of Justice (2016)\n",
      "4. Man of Steel (2013)\n",
      "5. The Avengers (2012)\n",
      "6. Iron Man (2008)\n",
      "7. Captain Amer...\n",
      "\n",
      "\u001b[0m1. Batman Begins (2005): A young Bruce Wayne becomes Batman to protect Gotham City from corruption and crime, facing his fears and training under the guidance of Ra's al Ghul.\n",
      "2. The Dark Knight Rises (2012): Batman returns from exile to save Gotham City from the ruthless terrorist Bane, who plans to destroy it, while facing personal sacrifices and a new ally, Catwoman.\n",
      "3. Batman v Superman: Dawn of Justice (2016): Batman and Superman clash as their ideologies collide, leading to an epic battle, while a new threat, Lex Luthor, manipulates them into a confrontation.\n",
      "4. Man of Steel (2013): The origin story of Superman, as he embraces his powers and faces General Zod, a fellow Kryptonian, who threatens Earth's existence.\n",
      "5. The Avengers (2012): Earth's mightiest heroes, including Iron Man, Captain America, Thor, and Hulk, unite to stop Loki and his alien army from conquering the world.\n",
      "6. Iron Man (2008): Billionaire Tony Stark builds a high-tech suit to escape captivity and becomes the superhero Iron Man, using his technology to fight against evil forces.\n",
      "7. Captain America: The Winter Soldier (2014): Captain America teams up with Black Widow and Falcon to uncover a conspiracy within S.H.I.E.L.D., facing a powerful assassin known as the Winter Soldier.\n",
      "8. The Amazing Spider-Man (2012): Peter Parker, a high school student, gains spider-like abilities and becomes Spider-Man, battling against the Lizard and discovering his father's secrets.\n",
      "9. Watchmen (2009): Set in an alternate reality, a group of retired vigilantes investigates the murder of one of their own, uncovering a plot that could lead to global destruction.\n",
      "10. Sin City (2005): A neo-noir anthology film, showcasing various interconnected stories set in the crime-ridden and corrupt Basin City.\n",
      "11. V for Vendetta (2005): In a dystopian future, a masked vigilante known as V fights against a totalitarian regime, inspiring the people to rise up against oppression.\n",
      "12. Blade Runner 2049 (2017): A young blade runner uncovers a long-buried secret that leads him to seek out former blade runner Rick Deckard, while unraveling the mysteries of a future society.\n",
      "13. Inception (2010): A skilled thief enters people's dreams to steal information, but is tasked with planting an idea instead, leading to a complex and mind-bending journey.\n",
      "14. The Matrix (1999): A computer hacker discovers the truth about reality, joining a group of rebels to fight against sentient machines that have enslaved humanity in a simulated world.\n",
      "15. The Crow (1994): A musician, resurrected by a supernatural crow, seeks revenge against those who murdered him and his fiancée, unleashing his dark and vengeful powers."
     ]
    }
   ],
   "source": [
    "output = p.run(movie_name=\"The Dark Knight\")\n",
    "for o in output:\n",
    "    print(o.delta, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf55168e-6d82-41e9-9e98-51025dc0ed3d",
   "metadata": {},
   "source": [
    "**2. Feed streaming output to output parser**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61752c36-8e7b-462a-ab23-4e3453d4b0b0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;155;135;227m> Running module 5cfd9352-07a6-4edd-90ac-f60cf0727a31 with input: \n",
      "movie_name: Toy Story\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module 1ccba87d-4d06-4bc2-bd7f-f044059a4091 with input: \n",
      "messages: Please generate related movies to Toy Story. Output with the following JSON format: \n",
      "\n",
      "\n",
      "\n",
      "Here's a JSON schema to follow:\n",
      "{\"title\": \"Movies\", \"description\": \"Object representing a list of movies.\", \"typ...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module 4ea05b9b-e4e2-4831-92d0-56790038c551 with input: \n",
      "input: <generator object llm_chat_callback.<locals>.wrap.<locals>.wrapped_llm_chat.<locals>.wrapped_gen at 0x2978e7760>\n",
      "\n",
      "\u001b[0mmovies=[Movie(name='Finding Nemo', year=2003), Movie(name='Cars', year=2006), Movie(name='Ratatouille', year=2007), Movie(name='WALL-E', year=2008), Movie(name='Up', year=2009), Movie(name='Inside Out', year=2015), Movie(name='Coco', year=2017), Movie(name='Incredibles 2', year=2018), Movie(name='Toy Story 4', year=2019), Movie(name='Onward', year=2020)]\n"
     ]
    }
   ],
   "source": [
    "p = QueryPipeline(\n",
    "    chain=[\n",
    "        json_prompt_tmpl,\n",
    "        llm.as_query_component(streaming=True),\n",
    "        output_parser,\n",
    "    ],\n",
    "    verbose=True,\n",
    ")\n",
    "output = p.run(movie_name=\"Toy Story\")\n",
    "print(output)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4572cdd4-2c94-4871-a496-623e9779e0db",
   "metadata": {},
   "source": [
    "## Chain Together Query Rewriting Workflow (prompts + LLM) with Retrieval\n",
    "\n",
    "Here we try a slightly more complex workflow where we send the input through two prompts before initiating retrieval.\n",
    "\n",
    "1. Generate question about given topic.\n",
    "2. Hallucinate answer given question, for better retrieval.\n",
    "\n",
    "Since each prompt only takes in one input, note that the `QueryPipeline` will automatically chain LLM outputs into the prompt and then into the LLM. \n",
    "\n",
    "You'll see how to define links more explicitly in the next section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c6908b1-5819-4f34-a06c-2f8b9fc81c3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.postprocessor import CohereRerank\n",
    "\n",
    "# generate question regarding topic\n",
    "prompt_str1 = \"Please generate a concise question about Paul Graham's life regarding the following topic {topic}\"\n",
    "prompt_tmpl1 = PromptTemplate(prompt_str1)\n",
    "# use HyDE to hallucinate answer.\n",
    "prompt_str2 = (\n",
    "    \"Please write a passage to answer the question\\n\"\n",
    "    \"Try to include as many key details as possible.\\n\"\n",
    "    \"\\n\"\n",
    "    \"\\n\"\n",
    "    \"{query_str}\\n\"\n",
    "    \"\\n\"\n",
    "    \"\\n\"\n",
    "    'Passage:\"\"\"\\n'\n",
    ")\n",
    "prompt_tmpl2 = PromptTemplate(prompt_str2)\n",
    "\n",
    "llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
    "retriever = index.as_retriever(similarity_top_k=5)\n",
    "p = QueryPipeline(\n",
    "    chain=[prompt_tmpl1, llm, prompt_tmpl2, llm, retriever], verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03eb2534-a69a-46fd-b539-84ae93e2e5bb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;155;135;227m> Running module 5b44aee5-2afe-4adb-8c56-370b047e5a8f with input: \n",
      "topic: college\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module 945590ad-31c1-4e26-b319-c43c887dd0c2 with input: \n",
      "messages: Please generate a concise question about Paul Graham's life regarding the following topic college\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module 0268b84c-3bed-46f2-9860-484763ec992c with input: \n",
      "query_str: assistant: How did Paul Graham's college experience shape his career and entrepreneurial mindset?\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module f3cbf237-e973-4e8a-ae8b-35fa7527c1b6 with input: \n",
      "messages: Please write a passage to answer the question\n",
      "Try to include as many key details as possible.\n",
      "\n",
      "\n",
      "assistant: How did Paul Graham's college experience shape his career and entrepreneurial mindset?\n",
      "\n",
      "\n",
      "Pass...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module da0f3129-7404-4c96-a8c8-35c01af710d8 with input: \n",
      "input: assistant: Paul Graham's college experience played a pivotal role in shaping his career and entrepreneurial mindset. As a student at Cornell University, Graham immersed himself in the world of compute...\n",
      "\n",
      "\u001b[0m"
     ]
    },
    {
     "data": {
      "text/plain": [
       "5"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nodes = p.run(topic=\"college\")\n",
    "len(nodes)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c073634-55d4-4066-977f-fbc189089b95",
   "metadata": {},
   "source": [
    "## Create a Full RAG Pipeline as a DAG\n",
    "\n",
    "Here we chain together a full RAG pipeline consisting of query rewriting, retrieval, reranking, and response synthesis.\n",
    "\n",
    "Here we can't use `chain` syntax because certain modules depend on multiple inputs (for instance, response synthesis expects both the retrieved nodes and the original question). Instead we'll construct a DAG explicitly, through `add_modules` and then `add_link`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66625e37-a0f2-45f4-afaa-e93948444e97",
   "metadata": {},
   "source": [
    "### 1. RAG Pipeline with Query Rewriting\n",
    "\n",
    "We use an LLM to rewrite the query first before passing it to our downstream modules - retrieval/reranking/synthesis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "531677b6-0e6a-4002-9fdd-2096f1a83685",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.postprocessor import CohereRerank\n",
    "from llama_index.response_synthesizers import TreeSummarize\n",
    "from llama_index import ServiceContext\n",
    "\n",
    "# define modules\n",
    "prompt_str = \"Please generate a question about Paul Graham's life regarding the following topic {topic}\"\n",
    "prompt_tmpl = PromptTemplate(prompt_str)\n",
    "llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
    "retriever = index.as_retriever(similarity_top_k=3)\n",
    "reranker = CohereRerank()\n",
    "summarizer = TreeSummarize(\n",
    "    service_context=ServiceContext.from_defaults(llm=llm)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ea2bd1e-8c32-420b-b557-eb7dbe81718f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# define query pipeline\n",
    "p = QueryPipeline(verbose=True)\n",
    "p.add_modules(\n",
    "    {\n",
    "        \"llm\": llm,\n",
    "        \"prompt_tmpl\": prompt_tmpl,\n",
    "        \"retriever\": retriever,\n",
    "        \"summarizer\": summarizer,\n",
    "        \"reranker\": reranker,\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e4e2af5-f7c2-4ab8-bb7c-76c369b63650",
   "metadata": {},
   "source": [
    "Next we draw links between modules with `add_link`. `add_link` takes in the source/destination module ids, and optionally the `source_key` and `dest_key`. Specify the `source_key` or `dest_key` if there are multiple outputs/inputs respectively.\n",
    "\n",
    "You can view the set of input/output keys for each module through `module.as_query_component().input_keys` and `module.as_query_component().output_keys`. \n",
    "\n",
    "Here we explicitly specify `dest_key` for the `reranker` and `summarizer` modules because they take in two inputs (query_str and nodes). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f0cccaf-8887-48f7-96b6-1d5458987022",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "required_keys={'query_str', 'nodes'} optional_keys=set()\n"
     ]
    }
   ],
   "source": [
    "p.add_link(\"prompt_tmpl\", \"llm\")\n",
    "p.add_link(\"llm\", \"retriever\")\n",
    "p.add_link(\"retriever\", \"reranker\", dest_key=\"nodes\")\n",
    "p.add_link(\"llm\", \"reranker\", dest_key=\"query_str\")\n",
    "p.add_link(\"reranker\", \"summarizer\", dest_key=\"nodes\")\n",
    "p.add_link(\"llm\", \"summarizer\", dest_key=\"query_str\")\n",
    "\n",
    "# look at summarizer input keys\n",
    "print(summarizer.as_query_component().input_keys)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cf69738-b409-408a-ac01-62bd4c2c2db4",
   "metadata": {},
   "source": [
    "We use `networkx` to store the graph representation. This gives us an easy way to view the DAG! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ef45868-ce07-4190-b12e-a70b287a491f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rag_dag.html\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "        <iframe\n",
       "            width=\"100%\"\n",
       "            height=\"600px\"\n",
       "            src=\"rag_dag.html\"\n",
       "            frameborder=\"0\"\n",
       "            allowfullscreen\n",
       "            \n",
       "        ></iframe>\n",
       "        "
      ],
      "text/plain": [
       "<IPython.lib.display.IFrame at 0x2b44698a0>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## create graph\n",
    "from pyvis.network import Network\n",
    "\n",
    "net = Network(notebook=True, cdn_resources=\"in_line\", directed=True)\n",
    "net.from_nx(p.dag)\n",
    "net.show(\"rag_dag.html\")\n",
    "\n",
    "## another option using `pygraphviz`\n",
    "# from networkx.drawing.nx_agraph import to_agraph\n",
    "# from IPython.display import Image\n",
    "# agraph = to_agraph(p.dag)\n",
    "# agraph.layout(prog=\"dot\")\n",
    "# agraph.draw('rag_dag.png')\n",
    "# display(Image('rag_dag.png'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "574ea13f-88e2-4393-8bd3-97fe436ffc0c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;155;135;227m> Running module prompt_tmpl with input: \n",
      "topic: YC\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module llm with input: \n",
      "messages: Please generate a question about Paul Graham's life regarding the following topic YC\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module retriever with input: \n",
      "input: assistant: What role did Paul Graham play in the founding and development of Y Combinator (YC)?\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module reranker with input: \n",
      "query_str: assistant: What role did Paul Graham play in the founding and development of Y Combinator (YC)?\n",
      "nodes: [NodeWithScore(node=TextNode(id_='543f958b-2c46-4c0f-b046-22e0a60ea950', embedding=None, metadata={'file_path': '../data/paul_graham/paul_graham_essay.txt', 'file_name': 'paul_graham_essay.txt', 'file...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module summarizer with input: \n",
      "query_str: assistant: What role did Paul Graham play in the founding and development of Y Combinator (YC)?\n",
      "nodes: [NodeWithScore(node=TextNode(id_='6b43fef2-821a-49a2-a043-9d816b70560f', embedding=None, metadata={'file_path': '../data/paul_graham/paul_graham_essay.txt', 'file_name': 'paul_graham_essay.txt', 'file...\n",
      "\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "response = p.run(topic=\"YC\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "998375e4-a8b4-4229-b034-6280d61e69f2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Paul Graham played a significant role in the founding and development of Y Combinator (YC). He was one of the co-founders of YC and was actively involved in its operations. He initially intended YC to be a part-time job, but as it grew, he became more excited about it and it started to take up a larger portion of his attention. Paul Graham worked on various aspects of YC, including selecting and helping founders, writing essays, and working on YC's internal software. However, over time, he gradually reduced his involvement in certain areas, such as working on the programming language Arc, as YC's infrastructure became more dependent on it. Eventually, Paul Graham decided to hand over the reins of YC to someone else and recruited Sam Altman to take over as president.\n"
     ]
    }
   ],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d976a65-1061-4b4a-8fdd-5dcf4c1bece9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;155;135;227m> Running module prompt_tmpl with input: \n",
      "topic: YC\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module llm with input: \n",
      "messages: Please generate a question about Paul Graham's life regarding the following topic YC\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module retriever with input: \n",
      "input: assistant: What role did Paul Graham play in the founding and development of Y Combinator (YC)?\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module reranker with input: \n",
      "query_str: assistant: What role did Paul Graham play in the founding and development of Y Combinator (YC)?\n",
      "nodes: [NodeWithScore(node=TextNode(id_='543f958b-2c46-4c0f-b046-22e0a60ea950', embedding=None, metadata={'file_path': '../data/paul_graham/paul_graham_essay.txt', 'file_name': 'paul_graham_essay.txt', 'file...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module summarizer with input: \n",
      "query_str: assistant: What role did Paul Graham play in the founding and development of Y Combinator (YC)?\n",
      "nodes: [NodeWithScore(node=TextNode(id_='6b43fef2-821a-49a2-a043-9d816b70560f', embedding=None, metadata={'file_path': '../data/paul_graham/paul_graham_essay.txt', 'file_name': 'paul_graham_essay.txt', 'file...\n",
      "\n",
      "\u001b[0mPaul Graham played a significant role in the founding and development of Y Combinator (YC). He was one of the co-founders of YC and was actively involved in its early stages. He helped establish the Summer Founders Program (SFP) and was responsible for selecting and funding the initial batch of startups. Graham also played a key role in shaping the funding model for YC, based on previous deals and experiences. As YC grew, Graham's involvement shifted, and he focused more on writing essays and working on YC, while also being involved in other projects such as Hacker News. However, he eventually decided to hand over the reins of YC to Sam Altman and retire from his role in the organization.\n"
     ]
    }
   ],
   "source": [
    "# you can do async too\n",
    "response = await p.arun(topic=\"YC\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce083e06-7943-4bc9-b134-ce4b7aa11f6a",
   "metadata": {},
   "source": [
    "### 2. RAG Pipeline without Query Rewriting\n",
    "\n",
    "Here we setup a RAG pipeline without the query rewriting step. \n",
    "\n",
    "Here we need a way to link the input query to both the retriever, reranker, and summarizer. We can do this by defining a special `InputComponent`, allowing us to link the inputs to multiple downstream modules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab324158-1491-46a7-95fd-c80ecd08957f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.postprocessor import CohereRerank\n",
    "from llama_index.response_synthesizers import TreeSummarize\n",
    "from llama_index import ServiceContext\n",
    "from llama_index.query_pipeline import InputComponent\n",
    "\n",
    "retriever = index.as_retriever(similarity_top_k=5)\n",
    "summarizer = TreeSummarize(\n",
    "    service_context=ServiceContext.from_defaults(\n",
    "        llm=OpenAI(model=\"gpt-3.5-turbo\")\n",
    "    )\n",
    ")\n",
    "reranker = CohereRerank()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "230cc32f-f155-4b72-8407-7618aa07622e",
   "metadata": {},
   "outputs": [],
   "source": [
    "p = QueryPipeline(verbose=True)\n",
    "p.add_modules(\n",
    "    {\n",
    "        \"input\": InputComponent(),\n",
    "        \"retriever\": retriever,\n",
    "        \"summarizer\": summarizer,\n",
    "    }\n",
    ")\n",
    "p.add_link(\"input\", \"retriever\")\n",
    "p.add_link(\"input\", \"summarizer\", dest_key=\"query_str\")\n",
    "p.add_link(\"retriever\", \"summarizer\", dest_key=\"nodes\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a64384a-6574-464a-a2a4-fa5f25091025",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;155;135;227m> Running module input with input: \n",
      "input: what did the author do in YC\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module retriever with input: \n",
      "input: what did the author do in YC\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module summarizer with input: \n",
      "query_str: what did the author do in YC\n",
      "nodes: [NodeWithScore(node=TextNode(id_='d2fc236f-2120-4f43-92d9-6c8e8725b806', embedding=None, metadata={'file_path': '../data/paul_graham/paul_graham_essay.txt', 'file_name': 'paul_graham_essay.txt', 'file...\n",
      "\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "output = p.run(input=\"what did the author do in YC\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82f261e6-d4f9-444e-904c-c253f97123d2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The author had a diverse range of responsibilities at YC, including writing essays, working on YC's internal software, funding and supporting startups, dealing with disputes between cofounders, identifying dishonesty, and addressing issues with people mistreating startups. They were also involved in the Summer Founders Program and the batch model of funding startups. Additionally, the author mentioned that they wrote all of YC's internal software in Arc, although they gradually stopped working on it.\n"
     ]
    }
   ],
   "source": [
    "print(str(output))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76c11f28-983d-4b3e-94a1-541c2804b989",
   "metadata": {},
   "source": [
    "## Defining a Custom Component in a Query Pipeline\n",
    "\n",
    "You can easily define a custom component. Simply subclass a `QueryComponent`, implement validation/run functions + some helpers, and plug it in.\n",
    "\n",
    "Let's wrap the related movie generation prompt+LLM chain from the first example into a custom component."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c470d37-a427-44a0-9b66-c83cbcac2545",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.query_pipeline import (\n",
    "    CustomQueryComponent,\n",
    "    InputKeys,\n",
    "    OutputKeys,\n",
    ")\n",
    "from typing import Dict, Any\n",
    "from llama_index.llms.llm import BaseLLM\n",
    "from pydantic import Field\n",
    "\n",
    "\n",
    "class RelatedMovieComponent(CustomQueryComponent):\n",
    "    \"\"\"Related movie component.\"\"\"\n",
    "\n",
    "    llm: BaseLLM = Field(..., description=\"OpenAI LLM\")\n",
    "\n",
    "    def _validate_component_inputs(\n",
    "        self, input: Dict[str, Any]\n",
    "    ) -> Dict[str, Any]:\n",
    "        \"\"\"Validate component inputs during run_component.\"\"\"\n",
    "        # NOTE: this is OPTIONAL but we show you here how to do validation as an example\n",
    "        return input\n",
    "\n",
    "    @property\n",
    "    def _input_keys(self) -> set:\n",
    "        \"\"\"Input keys dict.\"\"\"\n",
    "        # NOTE: These are required inputs. If you have optional inputs please override\n",
    "        # `optional_input_keys_dict`\n",
    "        return {\"movie\"}\n",
    "\n",
    "    @property\n",
    "    def _output_keys(self) -> set:\n",
    "        return {\"output\"}\n",
    "\n",
    "    def _run_component(self, **kwargs) -> Dict[str, Any]:\n",
    "        \"\"\"Run the component.\"\"\"\n",
    "        # use QueryPipeline itself here for convenience\n",
    "        prompt_str = \"Please generate related movies to {movie_name}\"\n",
    "        prompt_tmpl = PromptTemplate(prompt_str)\n",
    "        p = QueryPipeline(chain=[prompt_tmpl, llm])\n",
    "        return {\"output\": p.run(movie_name=kwargs[\"movie\"])}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69dc1a62-1e84-45c5-9f12-cac7778bd46f",
   "metadata": {},
   "source": [
    "Let's try the custom component out! We'll also add a step to convert the output to Shakespeare."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8ff382d-baf8-46e2-b4c0-477a07a41219",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
    "component = RelatedMovieComponent(llm=llm)\n",
    "\n",
    "# let's add some subsequent prompts for fun\n",
    "prompt_str = \"\"\"\\\n",
    "Here's some text:\n",
    "\n",
    "{text}\n",
    "\n",
    "Can you rewrite this in the voice of Shakespeare?\n",
    "\"\"\"\n",
    "prompt_tmpl = PromptTemplate(prompt_str)\n",
    "\n",
    "p = QueryPipeline(chain=[component, prompt_tmpl, llm], verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7f6aa04-4443-4511-966d-c71cbd268efa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1;3;38;2;155;135;227m> Running module 4987ea07-87d3-43ec-98dc-2bec13b2a237 with input: \n",
      "movie: Love Actually\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module 787b9daf-ee18-4e21-85ed-d04fe6dcc347 with input: \n",
      "text: assistant: 1. \"Valentine's Day\" (2010) - A star-studded ensemble cast explores interconnected love stories on Valentine's Day in Los Angeles.\n",
      "\n",
      "2. \"New Year's Eve\" (2011) - Similar to \"Love Actually,\" ...\n",
      "\n",
      "\u001b[0m\u001b[1;3;38;2;155;135;227m> Running module cd470014-5b67-4e21-b769-bb5f3aa9f2b2 with input: \n",
      "messages: Here's some text:\n",
      "\n",
      "assistant: 1. \"Valentine's Day\" (2010) - A star-studded ensemble cast explores interconnected love stories on Valentine's Day in Los Angeles.\n",
      "\n",
      "2. \"New Year's Eve\" (2011) - Similar t...\n",
      "\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "output = p.run(movie=\"Love Actually\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "726ee679-e6a3-4e57-9566-00ef9f40b837",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "assistant: 1. \"Valentine's Daye\" (2010) - A troupe of stars doth explore interconnected love stories on Valentine's Daye in fair Los Angeles.\n",
      "\n",
      "2. \"New Year's Eve\" (2011) - Similar to \"Love Actually,\" this play doth followeth multiple characters as they navigate love and relationships on New Year's Eve in fair New York City.\n",
      "\n",
      "3. \"He's Just Not That Into Thee\" (2009) - This romantic comedy doth feature intersecting storylines that explore the complexities of modern relationships and the search for true love.\n",
      "\n",
      "4. \"Crazy, Stupid, Love\" (2011) - A middle-aged man's life doth unravel when his wife doth ask him for a divorce, leading him to seeketh guidance from a young bachelor who doth help him rediscover love.\n",
      "\n",
      "5. \"The Holiday\" (2006) - Two women, one from fair Los Angeles and the other from England, doth swap homes during the Christmas season and unexpectedly findeth love in their new surroundings.\n",
      "\n",
      "6. \"Four Weddings and a Funeral\" (1994) - This British romantic comedy doth followeth a group of friends as they attendeth various weddings and a funeral, exploring the ups and downs of love and relationships.\n",
      "\n",
      "7. \"Notting Hill\" (1999) - A famous actress and a humble bookstore owner doth fall in love in the vibrant neighborhood of Notting Hill in London, leading to a series of comedic and heartfelt moments.\n",
      "\n",
      "8. \"Bridget Jones's Diary\" (2001) - Bridget Jones, a quirky and relatable woman, doth navigate her love life whilst documenting her experiences in a diary, leading to unexpected romantic entanglements.\n",
      "\n",
      "9. \"About Time\" (2013) - A young man doth discover he hath the ability to time travel and useth it to find love, but soon realizeth that altering the past can have unforeseen consequences.\n",
      "\n",
      "10. \"The Best Exotic Marigold Hotel\" (2011) - A group of British retirees doth decide to spendeth their golden years in a seemingly luxurious hotel in India, where they findeth love, friendship, and new beginnings.\n"
     ]
    }
   ],
   "source": [
    "print(str(output))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama_index_v2",
   "language": "python",
   "name": "llama_index_v2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
